# -*- coding: UTF-8 -*-
import os
import rospy
import numpy as np
import matplotlib.pyplot as plt
from sensor_msgs.msg import Imu
import threading

def NORM(x):
    return np.math.sqrt(np.dot(x,x))

def dis_ACC(x, y):
    scale = 0.5
    diff = x - y
    dis = np.dot(x, y) / (NORM(x) * NORM(y) + 1e-8)
    dis = (1 - scale * dis) * NORM(diff)
    return dis


def estimate_dtw(A, B, dis_func=dis_ACC):
    N_A = len(A)
    N_B = len(B)

    D = np.zeros([N_A, N_B])
    D[0, 0] = dis_func(A[0], B[0])

    #左边一列
    for i in range(1, N_A):
        D[i, 0] = D[i - 1, 0] + dis_func(A[i], B[0])
    #下边一行
    for j in range(1, N_B):
        D[0, j] = D[0, j - 1] + dis_func(A[0], B[j])
    #中间部分
    for i in range(1, N_A):
        for j in range(1, N_B):
            D[i, j] = dis_func(A[i], B[j]) + min(D[i - 1, j], D[i, j - 1], D[i - 1, j - 1])

    # 路径回溯
    i = N_A - 1
    j = N_B - 1
    count = 0
    d = np.zeros(max(N_A, N_B) * 3)
    path = []
    while True:
        if i > 0 and j > 0:
            path.append((i, j))
            m = min(D[i - 1, j], D[i, j - 1], D[i - 1, j - 1])
            if m == D[i - 1, j - 1]:
                d[count] = D[i, j] - D[i - 1, j - 1]
                i = i - 1
                j = j - 1
                count = count + 1

            elif m == D[i, j - 1]:
                d[count] = D[i, j] - D[i, j - 1]
                j = j - 1
                count = count + 1

            elif m == D[i - 1, j]:
                d[count] = D[i, j] - D[i - 1, j]
                i = i - 1
                count = count + 1

        elif i == 0 and j == 0:
            path.append((i, j))
            d[count] = D[i, j]
            count = count + 1
            break

        elif i == 0:
            path.append((i, j))
            d[count] = D[i, j] - D[i, j - 1]
            j = j - 1
            count = count + 1

        elif j == 0:
            path.append((i, j))
            d[count] = D[i, j] - D[i - 1, j]
            i = i - 1
            count = count + 1

    mean = np.sum(d) / count
    return mean, path[::-1], D


def get_mov_list(path):
    list_files = []
    list_labs = []

    if os.path.exists(path):

        # 遍历文件夹，找到.npy文件
        for root, ds, fs in os.walk(path):
            for f in fs:
                if f.endswith('.npy'):
                    fullname = os.path.join(root, f)
                    list_files.append(fullname)

        for file in list_files:
            lab = os.path.split(os.path.split(file)[0])[-1]
            list_labs.append(lab)

    else:
        print('this path not exist')

    return list_files, list_labs

def imuVisualize(data):
    fig = plt.figure()
    ax = plt.axes()
    i=0
    xlist=[]
    ylist=[]
    zlist=[]
    index=[]
    for item in data:
        index.append(i)
        i+=1
        xlist.append(item[0])
        ylist.append(item[1])
        zlist.append(item[2])
    plt.plot(index, xlist,color='r')
    plt.plot(index, ylist, color='g')
    plt.plot(index, zlist, color='b')

datalist=[]

imu_buf=threading.Lock()

def handleIMU(data):
    imudata=[]
    imu_buf.acquire()
    imudata.append(data.linear_acceleration.x)
    imudata.append(data.linear_acceleration.y)
    imudata.append(data.linear_acceleration.z)
    datalist.append(imudata)
    imu_buf.release()

if __name__ == "__main__":

    rospy.init_node('IMUhandle',anonymous=True)
    imusub = rospy.Subscriber("/camera/imu", Imu,handleIMU, queue_size=20)
    i=0

    mydata,labs=get_mov_list("../imu_template")
    mydata1=np.load(mydata.pop())
    mydata2=np.load(mydata.pop())
    imuVisualize(mydata1)
    imuVisualize(mydata2)
    score=estimate_dtw(mydata1,mydata2)
    print(score[0])
    plt.show()
    #
    # while True:
    #     if len(datalist)>=300:
    #         imu_buf.acquire()
    #         check=[]
    #         check=datalist
    #         datalist=[]
    #         imu_buf.release()
    #         #path = os.path.join("imu_record",str(i))
    #         # os.makedirs(path)
    #         # file_name = os.path.join(path, str(i) + ".npy")
    #         # np.save(file_name, np.array(check))
    #         i+=1
    #         #i=imuVisualize2(check,i)
    #         scores=estimate_dtw(check,mydata)
    #         check = []
    #         print("done with "+str(i))
    #         print("DTW score: "+str(scores[0]))

    # files, labs = get_mov_list("imu_template")
    # for file in files:
    #     # 加载动作数据
    #     data_train = np.load(file)
    #     imuVisualize(data_train)
    #     # 计算测试动作与存储动作之间的距离
    #     print("find data!")
    # plt.show()
        #score, _, _ = estimate_dtw(np.array(test_mov), data_train)
        #scores.append(score)

